#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import rclpy
from rclpy.node import Node
import moveit_commander
import geometry_msgs.msg
from moveit_commander import RobotTrajectory
from trajectory_msgs.msg import JointTrajectoryPoint
import time

def scale_trajectory_speed(traj, scale):
    # 创建新的轨迹对象并初始化
    new_traj = RobotTrajectory()
    new_traj.joint_trajectory = traj.joint_trajectory

    # 获取关节和轨迹点数量
    n_joints = len(traj.joint_trajectory.joint_names)
    n_points = len(traj.joint_trajectory.points)
    points = list(traj.joint_trajectory.points)
    
    # 调整时间、速度和加速度
    for i in range(n_points):
        point = JointTrajectoryPoint()
        point.positions = traj.joint_trajectory.points[i].positions
        point.time_from_start = traj.joint_trajectory.points[i].time_from_start / scale
        point.velocities = list(traj.joint_trajectory.points[i].velocities)
        point.accelerations = list(traj.joint_trajectory.points[i].accelerations)
        
        for j in range(n_joints):
            point.velocities[j] *= scale
            point.accelerations[j] *= scale * scale

        points[i] = point

    new_traj.joint_trajectory.points = points
    return new_traj

def execute_trajectory(waypoints, speed):
        # 初始化 moveit_commander 和 rospy
        # moveit_commander.roscpp_initialize(sys.argv)
        arm_group = moveit_commander.MoveGroupCommander("manipulator")
        # 规划并执行笛卡尔路径
        (plan, fraction) = arm_group.compute_cartesian_path(waypoints, 0.01, False)
        print(f"Fraction of path achieved: {fraction}")
        
        # 调整轨迹速度
        if speed == 'fast':
            scaled_plan = scale_trajectory_speed(plan, 1.5)
        elif speed == 'medium':            scaled_plan = scale_trajectory_speed(plan, 1.25)
        elif speed == 'slow':
            scaled_plan = scale_trajectory_speed(plan, 1)
        
        # 执行轨迹
        success = arm_group.execute(scaled_plan, wait=True)
        arm_group.stop()
        arm_group.clear_pose_targets()
        if not success:
            print("轨迹执行失败")
            return
        
        rospy.sleep(3)
        arm_group.set_named_target("middle")
        arm_group.go(wait=True)

def trajectory_1():
    waypoints = []

    # 第1个目标位姿 30/-70/90/-20/0
    pose1 = geometry_msgs.msg.Pose()
    pose1.position.x = -0.0011310404075417224
    pose1.position.y = 0.06907060201161805
    pose1.position.z = 0.3073585902967819
    pose1.orientation.x = 2.688964454397768e-05
    pose1.orientation.y = -9.329175674129769e-05
    pose1.orientation.z = 0.2588310042911079
    pose1.orientation.w = 0.9659226168753127
    waypoints.append(pose1)

    # 第2个目标位姿 35/-55/110/-10/0
    pose2 = geometry_msgs.msg.Pose()
    pose2.position.x = 0.009665201991908984
    pose2.position.y = 0.09130426767319212
    pose2.position.z = 0.11498378100534505
    pose2.orientation.x = -0.11512122029744272
    pose2.orientation.y = 0.36498810976129026
    pose2.orientation.z = 0.27778707139570713
    pose2.orientation.w = 0.8811158421771251
    waypoints.append(pose2)

    # 4 35/-60/115/-10/0 
    pose4 = geometry_msgs.msg.Pose()
    pose4.position.x = 0.002689802791981763
    pose4.position.y = 0.0864335660110534
    pose4.position.z = 0.1016529166473485
    pose4.orientation.x = -0.11515741145952224
    pose4.orientation.y = 0.3650111716795012
    pose4.orientation.z = 0.27781956196470253
    pose4.orientation.w = 0.8810913154292472
    waypoints.append(pose4)

    # 5 35/-55/125/-30/0
    pose5 = geometry_msgs.msg.Pose()
    pose5.position.x = 0.040521254025108394
    pose5.position.y = 0.11294248356505611
    pose5.position.z = 0.08770454646571219
    pose5.orientation.x = -0.10292980365209553
    pose5.orientation.y = 0.3261289152765809
    pose5.orientation.z = 0.2825969968796311
    pose5.orientation.w = 0.8962055140955537
    waypoints.append(pose5)

    # 第6个目标位姿 35/-50/122/-32/0
    pose6 = geometry_msgs.msg.Pose()
    pose6.position.x = 0.05010657634938402
    pose6.position.y = 0.1196340400624652
    pose6.position.z = 0.0942325579027638
    pose6.orientation.x = -0.10295037417731849
    pose6.orientation.y = 0.32615192764866546
    pose6.orientation.z = 0.28254401566504306
    pose6.orientation.w = 0.8962114816044716
    waypoints.append(pose6)

    # 第7个目标位姿 32/-45/120/-35/0
    pose7 = geometry_msgs.msg.Pose()
    pose7.position.x = 0.06747123197664025
    pose7.position.y = 0.11753887303608636
    pose7.position.z = 0.09779214944642986
    pose7.orientation.x = -0.09434732313454879
    pose7.orientation.y = 0.3287544696671605
    pose7.orientation.z = 0.25884543602305454
    pose7.orientation.w = 0.9033372136368811
    waypoints.append(pose7)

    # 第8个目标位姿 32/-58/138/-60/0
    pose8 = geometry_msgs.msg.Pose()
    pose8.position.x = 0.0904501718573451
    pose8.position.y = 0.13184367830868512
    pose8.position.z = 0.11971295004270945
    pose8.orientation.x = -0.047902301026064985
    pose8.orientation.y = 0.16691573525552947
    pose8.orientation.z = 0.2711677829418468
    pose8.orientation.w = 0.9467378414191111
    waypoints.append(pose8)

    # 第9个目标位姿 -32/-70/150/-78/0 举平
    pose9 = geometry_msgs.msg.Pose()
    pose9.position.x = 0.09432316993420176
    pose9.position.y = 0.13433610299295476
    pose9.position.z = 0.16009946275824632
    pose9.orientation.x = -0.004805209403384029
    pose9.orientation.y = 0.016742670630141412
    pose9.orientation.z = 0.27542811607470175
    pose9.orientation.w = 0.9611638496210213
    waypoints.append(pose9)


    # 第10个目标位姿 -20/-70/150/-78/0 举平移
    pose10 = geometry_msgs.msg.Pose()
    pose10.position.x = 0.11755533690651693
    pose10.position.y = 0.08668575883939508
    pose10.position.z = 0.16011078889265554
    pose10.orientation.x = -0.00302493770022203
    pose10.orientation.y = 0.017144256975920424
    pose10.orientation.z = 0.17347487077387133
    pose10.orientation.w = 0.9846844131063731
    waypoints.append(pose10)

    # # 第11个目标位姿 
    pose11 = geometry_msgs.msg.Pose()
    pose11.position.x = 0.009070536838216042
    pose11.position.y = 0.04719269375993429
    pose11.position.z = 0.3072916016195043
    pose11.orientation.x = -1.1967649283791688e-06
    pose11.orientation.y = 1.7336762845100254e-05
    pose11.orientation.z = 0.1735087311185056
    pose11.orientation.w = 0.9848323308683823
    waypoints.append(pose11)

    # 第14个目标位姿 0/-70/90/-20 初始位置
    pose14 = geometry_msgs.msg.Pose()
    pose14.position.x = 0.017405918539271842
    pose14.position.y = -1.2171559032328237e-06
    pose14.position.z = 0.30727696933440546
    pose14.orientation.x = 1.829736299180823e-06
    pose14.orientation.y = 8.08447651258564e-05
    pose14.orientation.z = 6.071146226902279e-06
    pose14.orientation.w = 0.9999999967119587
    waypoints.append(pose14)

        # 第15个目标位姿 送食位置
    pose15 = geometry_msgs.msg.Pose()
    pose15.position.x = 0.3578591546734353
    pose15.position.y = -4.5874290626404854e-05
    pose15.position.z = 0.38613742059913403
    pose15.orientation.x = -4.9109689661127625e-06
    pose15.orientation.y = -0.06983318184000316
    pose15.orientation.z = -2.5349309214965662e-05
    pose15.orientation.w = 0.9975586830093738
    waypoints.append(pose15)
    return waypoints

def trajectory_2():
    waypoints = []

    # 第1个目标位姿 0/-55/90/20/0
    pose1 = geometry_msgs.msg.Pose()
    pose1.position.x = -0.017099604521405784
    pose1.position.y = -2.8315387428728567e-05
    pose1.position.z = 0.11438909247802503
    pose1.orientation.x = -3.453025354057534e-05
    pose1.orientation.y = 0.46175405203633646
    pose1.orientation.z = -7.30327633778306e-05
    pose1.orientation.w = 0.887007998217548
    waypoints.append(pose1)

    # 第2个目标位姿 0/-58/115/-3/0
    pose2 = geometry_msgs.msg.Pose()
    pose2.position.x = 0.004651029454941641
    pose2.position.y = -2.046464646042786e-06
    pose2.position.z = 0.07989258904926122
    pose2.orientation.x = -3.9551522751682376e-05
    pose2.orientation.y = 0.4523160885130879
    pose2.orientation.z = 6.4712743132175465e-06
    pose2.orientation.w = 0.8918576985517478
    waypoints.append(pose2)

    # 第3个目标位姿 0/-80/131/-11/0
    pose3 = geometry_msgs.msg.Pose()
    pose3.position.x = 0.016710764091943137
    pose3.position.y = 4.7481613370422367e-05
    pose3.position.z = 0.06725314521298217
    pose3.orientation.x = 2.484310933627909e-06
    pose3.orientation.y = 0.34210428530462056
    pose3.orientation.z = 0.00014130528033512744
    pose3.orientation.w = 0.9396619807158642
    waypoints.append(pose3)

    # 4 0/-90/142/-26/0
    pose4 = geometry_msgs.msg.Pose()
    pose4.position.x = 0.061313835276576675
    pose4.position.y = 3.4229349294604146e-05
    pose4.position.z = 0.09607582904327464
    pose4.orientation.x = 8.189114554395947e-06
    pose4.orientation.y = 0.19079166522506064
    pose4.orientation.z = 9.397613957952172e-05
    pose4.orientation.w = 0.9816305473891067
    waypoints.append(pose4)

    # 5 0/-90/146/-32/0
    pose5 = geometry_msgs.msg.Pose()
    pose5.position.x = 0.06469858300599692
    pose5.position.y = 6.72000076315633e-05
    pose5.position.z = 0.07814980502181258
    pose5.orientation.x = -3.84250569528825e-06
    pose5.orientation.y = 0.20806965601170563
    pose5.orientation.z = 0.00018064511789783727
    pose5.orientation.w = 0.978113994174374
    waypoints.append(pose5)

    # 第7个目标位姿 0/-90/155/-65/0 举平new
    pose7 = geometry_msgs.msg.Pose()
    pose7.position.x = 0.10862679372272958
    pose7.position.y = 7.322812585541887e-05
    pose7.position.z = 0.151673237887705
    pose7.orientation.x = 1.843833939518914e-06
    pose7.orientation.y = 0.0002691494452088297
    pose7.orientation.z = 0.00018408295291220934
    pose7.orientation.w = 0.9999999468343201
    waypoints.append(pose7)

    # 第8个目标位姿 -15/-90/155/-68/0
    pose8 = geometry_msgs.msg.Pose()
    pose8.position.x = 0.10150451167124033
    pose8.position.y = -0.059668682905588945
    pose8.position.z = 0.1647479924791181 #0.15
    pose8.orientation.x = -0.003413749283397351
    pose8.orientation.y = -0.02598662685755732
    pose8.orientation.z = -0.1307826110962234
    pose8.orientation.w = 0.9910645539899258
    waypoints.append(pose8)

    # 第10个目标位姿 
    pose10 = geometry_msgs.msg.Pose()
    pose10.position.x = 0.12119300384217499
    pose10.position.y = -0.06496064015101635
    pose10.position.z = 0.16648646107426138
    pose10.orientation.x = -0.003406584361246141
    pose10.orientation.y = -0.02590040792578116
    pose10.orientation.z = -0.130787996023333
    pose10.orientation.w = 0.9910661250130935
    waypoints.append(pose10)

    # 第11个目标位姿 -15/-75/135/-68/0
    pose11 = geometry_msgs.msg.Pose()
    pose11.position.x = 0.10021590416090176
    pose11.position.y = -0.059275851310574126
    pose11.position.z = 0.24717377948372102
    pose11.orientation.x = -0.009100789371875279
    pose11.orientation.y = -0.06914916681497274
    pose11.orientation.z = -0.1303713558152669
    pose11.orientation.w = 0.9890090383532868
    waypoints.append(pose11)

    # 第12个目标位姿 0/-75/135/-68/0
    pose12 = geometry_msgs.msg.Pose()
    pose12.position.x = 0.1080342010592359
    pose12.position.y = -6.057566946212703e-06
    pose12.position.z = 0.24717779672627277
    pose12.orientation.x = 9.786502587199877e-06
    pose12.orientation.y = -0.06972668749366641
    pose12.orientation.z = -2.98995522507148e-05
    pose12.orientation.w = 0.9975661321744048
    waypoints.append(pose12)

    # 第14个目标位姿 0/-70/90/-20 设定这个位置为初始位置
    pose14 = geometry_msgs.msg.Pose()
    pose14.position.x = 0.017405918539271842
    pose14.position.y = -1.2171559032328237e-06
    pose14.position.z = 0.30727696933440546
    pose14.orientation.x = 1.829736299180823e-06
    pose14.orientation.y = 8.08447651258564e-05
    pose14.orientation.z = 6.071146226902279e-06
    pose14.orientation.w = 0.9999999967119587
    waypoints.append(pose14)
 
        # 第15个目标位姿 送食位置
    pose15 = geometry_msgs.msg.Pose()
    pose15.position.x = 0.3578591546734353
    pose15.position.y = -4.5874290626404854e-05
    pose15.position.z = 0.38613742059913403
    pose15.orientation.x = -4.9109689661127625e-06
    pose15.orientation.y = -0.06983318184000316
    pose15.orientation.z = -2.5349309214965662e-05
    pose15.orientation.w = 0.9975586830093738
    waypoints.append(pose15)
    return waypoints

def trajectory_3():
    waypoints = []

    # 第1个目标位姿 -30/-70/90/-20/0
    pose1 = geometry_msgs.msg.Pose()
    pose1.position.x = -0.0011302048333630965
    pose1.position.y = -0.06905961719492143
    pose1.position.z = 0.3073550976191969
    pose1.orientation.x = -1.7316002977395957e-05
    pose1.orientation.y = -7.170005905241826e-05
    pose1.orientation.z = -0.2588197909540237
    pose1.orientation.w = 0.9659256236221158
    waypoints.append(pose1)

    # 第2个目标位姿 -35/-55/110/-10/0
    pose2 = geometry_msgs.msg.Pose()
    pose2.position.x = 0.009670549573942186
    pose2.position.y = -0.09129423315793453
    pose2.position.z = 0.11500108680844905
    pose2.orientation.x = 0.1150590441241054
    pose2.orientation.y = 0.3649594495490963
    pose2.orientation.z = -0.2777586943432337
    pose2.orientation.w = 0.881144780536553
    waypoints.append(pose2)

    # 4 -35/-60/115/-10/0 
    pose4 = geometry_msgs.msg.Pose()
    pose4.position.x = 0.002689802791981763
    pose4.position.y = -0.08642607255938041
    pose4.position.z = 0.10169023602364938
    pose4.orientation.x = 0.11506762902098444
    pose4.orientation.y = 0.3649312412859769
    pose4.orientation.z = -0.2777332746713525
    pose4.orientation.w = 0.8811633549037843
    waypoints.append(pose4)

    # 5 -35/-55/125/-30/0
    pose5 = geometry_msgs.msg.Pose()
    pose5.position.x = 0.04053697379986093
    pose5.position.y = -0.11290902734774333
    pose5.position.z = 0.0876834203451664
    pose5.orientation.x = 0.10283628094347873
    pose5.orientation.y = 0.32616481198444003
    pose5.orientation.z = -0.2825256471019672
    pose5.orientation.w = 0.8962256822221084
    waypoints.append(pose5)

    # 第6个目标位姿 -35/-50/122/-32/0
    pose6 = geometry_msgs.msg.Pose()
    pose6.position.x = 0.05010120615902845
    pose6.position.y = -0.11958529726583864
    pose6.position.z = 0.09421890307418768
    pose6.orientation.x = 0.1028380818468409
    pose6.orientation.y = 0.3262024609161537
    pose6.orientation.z = -0.28248616522541353
    pose6.orientation.w = 0.8962242185248891
    waypoints.append(pose6)

    # 第7个目标位姿 -32/-45/120/-35/0
    pose7 = geometry_msgs.msg.Pose()
    pose7.position.x = 0.06745494524328567
    pose7.position.y = -0.11761062296318618
    pose7.position.z = 0.09681908582144214
    pose7.orientation.x = 0.09426878785030171
    pose7.orientation.y = 0.32868644227651406
    pose7.orientation.z = -0.259031952013786
    pose7.orientation.w = 0.9033167031206555
    waypoints.append(pose7)

    # 第8个目标位姿 32/-58/138/-60/0
    pose8 = geometry_msgs.msg.Pose()
    pose8.position.x = 0.09038448867848176
    pose8.position.y = -0.1319360201480157
    pose8.position.z = 0.1197025357922735
    pose8.orientation.x = 0.04787093957528583
    pose8.orientation.y = 0.16692700118128137
    pose8.orientation.z = -0.2714553558001961
    pose8.orientation.w = 0.9466550265160975
    waypoints.append(pose8)

    # 第9个目标位姿 -32/-70/150/-78/0 举平
    pose9 = geometry_msgs.msg.Pose()
    pose9.position.x = 0.0942933762027911
    pose9.position.y = -0.13439538389383587
    pose9.position.z = 0.16010801190356336
    pose9.orientation.x = 0.004799855100987425
    pose9.orientation.y = 0.0167298817281641
    pose9.orientation.z = -0.27562936709118274
    pose9.orientation.w = 0.9611064064115302
    waypoints.append(pose9)


    # 第10个目标位姿 -20/-70/150/-78/0 举平移
    pose10 = geometry_msgs.msg.Pose()
    pose10.position.x = 0.11754935088930349
    pose10.position.y = -0.08673064743303559
    pose10.position.z = 0.1601478830358012
    pose10.orientation.x = 0.0030148054312412937
    pose10.orientation.y = 0.017086839017522924
    pose10.orientation.z = -0.17361459296382126
    pose10.orientation.w = 0.984660816723509
    waypoints.append(pose10)

    # # 第11个目标位姿 -20/-70/90/-20/0
    pose11 = geometry_msgs.msg.Pose()
    pose11.position.x = 0.009071977232847944
    pose11.position.y = -0.047260477105669295
    pose11.position.z = 0.3073157690252368
    pose11.orientation.x = 3.856803300718398e-06
    pose11.orientation.y = 1.1296182120014837e-05
    pose11.orientation.z = -0.17367575774823565
    pose11.orientation.w = 0.9848028894291982
    waypoints.append(pose11)

    # 第14个目标位姿 0/-70/90/-20 设定这个位置为初始位置
    pose14 = geometry_msgs.msg.Pose()
    pose14.position.x = 0.017405918539271842
    pose14.position.y = -1.2171559032328237e-06
    pose14.position.z = 0.30727696933440546
    pose14.orientation.x = 1.829736299180823e-06
    pose14.orientation.y = 8.08447651258564e-05
    pose14.orientation.z = 6.071146226902279e-06
    pose14.orientation.w = 0.9999999967119587
    waypoints.append(pose14)


        # 第15个目标位姿 送食位置
    pose15 = geometry_msgs.msg.Pose()
    pose15.position.x = 0.3578591546734353
    pose15.position.y = -4.5874290626404854e-05
    pose15.position.z = 0.38613742059913403
    pose15.orientation.x = -4.9109689661127625e-06
    pose15.orientation.y = -0.06983318184000316
    pose15.orientation.z = -2.5349309214965662e-05
    pose15.orientation.w = 0.9975586830093738
    waypoints.append(pose15)
    return waypoints

def run_selected_trajectory(traj_index, speed):
    if traj_index == 1:
        waypoints = trajectory_1()
    elif traj_index == 2:
        waypoints = trajectory_2()
    elif traj_index == 3:
        waypoints = trajectory_3()
    else:
        print("无效的轨迹编号")
        return

    if speed not in ['fast', 'medium', 'slow']:
        print("无效的速度。请选择 'fast', 'medium' 或 'slow'.")
        return
    execute_trajectory(waypoints, speed)

if __name__ == '__main__':
    try:
        # 示例: 执行第二条轨迹的中速
        run_selected_trajectory(1, 'fast')
        run_selected_trajectory(3, 'fast')
    except rospy.ROSInterruptException:
        pass
